#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
File: sqliteSave.py
Author: StellaJL
Date: 2025-09-16 14:57:12
Version：1.0.0
Description: 通用SQLite数据保存器类
"""

import aiosqlite
import os
import yaml
import sqlite3
from logger import get_logger


class SqliteSaver:
    """
    SQLite数据保存器类
    - 会根据传入的table_name创建表
    - 会根据传入的data动态创建表结构
    - 会根据传入的index_field创建索引
    """
    
    def __init__(self, config_path='config.yaml'):
        """
        初始化SQLite保存器
        
        Args:
            config_path: 配置文件路径
        """
        self.config_path = config_path
        self.logger = get_logger('SqliteSaver')
        self._config = None
        self._sqlite_enabled = False
        self._db_path = None
        self._database_name = None
        self._initialized = False
        # 索引缓存：记录已经创建过索引的表和字段组合
        # 格式: {table_name: {index_field: True}}
        self._created_indexes = {}
        
        # 初始化配置
        self._load_config()
    
    def _load_config(self):
        """加载配置文件"""
        try:
            if not os.path.exists(self.config_path):
                raise FileNotFoundError(f"配置文件 {self.config_path} 未找到。")
            
            with open(self.config_path, 'r', encoding='utf-8') as f:
                self._config = yaml.safe_load(f)
            
            # 检查SQLite配置
            if self._config and 'storage' in self._config:
                storage_config = self._config['storage']
                
                if ('sources' in storage_config and 
                    'sqlite' in storage_config['sources'] and 
                    storage_config['sources']['sqlite'].get('enabled', False)):
                    
                    sqlite_config = storage_config['sources']['sqlite']
                    sqlite_path = sqlite_config.get('path', './crawled_data/sqlite')
                    sqlite_filename = sqlite_config.get('filename', 'temp.db')
                    
                    # 确保SQLite目录存在
                    os.makedirs(sqlite_path, exist_ok=True)
                    
                    # 构建完整的数据库路径
                    self._db_path = os.path.join(sqlite_path, sqlite_filename)
                    self._database_name = sqlite_config.get('database', 'temp_spider')
                    self._sqlite_enabled = True
                    
                    self.logger.info(f"SQLite配置加载成功: {self._db_path}")
                else:
                    self.logger.warning("SQLite未启用或配置不完整")
            else:
                self.logger.error("没有找到存储配置")
                
            self._initialized = True
            
        except Exception as e:
            self.logger.error(f"加载配置失败: {e}")
            raise
    
    def _validate_index_field(self, index_field: str, data: list) -> bool:
        """
        验证索引字段是否有效
        
        Args:
            index_field: 要创建索引的字段名
            data: 数据列表
            
        Returns:
            bool: 字段是否有效
        """
        if not index_field or not isinstance(index_field, str):
            return False
        
        if not data or len(data) == 0:
            return False
        
        # 检查字段是否存在于数据中
        sample_item = data[0]
        return index_field in sample_item
    
    async def _create_index(self, db, table_name: str, index_field: str):
        """
        为指定字段创建索引
        
        Args:
            db: 数据库连接对象
            table_name: 表名
            index_field: 要创建索引的字段名
        """
        try:
            index_name = f"idx_{table_name}_{index_field}"
            create_index_sql = f"CREATE INDEX {index_name} ON {table_name} ({index_field})"
            
            await db.execute(create_index_sql)
            self.logger.info(f"成功为字段 '{index_field}' 创建索引: {index_name}")
            
            # 记录已创建的索引到缓存
            if table_name not in self._created_indexes:
                self._created_indexes[table_name] = {}
            self._created_indexes[table_name][index_field] = True
            
        except sqlite3.OperationalError as e:
            if "already exists" in str(e).lower():
                self.logger.warning(f"索引 {index_name} 已存在，跳过创建")
                # 即使索引已存在，也记录到缓存中
                if table_name not in self._created_indexes:
                    self._created_indexes[table_name] = {}
                self._created_indexes[table_name][index_field] = True
            else:
                self.logger.error(f"创建索引失败: {e}")
                raise
        except Exception as e:
            self.logger.error(f"创建索引时发生错误: {e}")
            raise
    
    async def _create_index_if_not_exists(self, db, table_name: str, index_field: str):
        """
        检查索引是否存在，如果不存在则创建（使用缓存优化）
        
        Args:
            db: 数据库连接对象
            table_name: 表名
            index_field: 要创建索引的字段名
        """
        try:
            # 首先检查缓存，避免重复查询数据库
            if (table_name in self._created_indexes and 
                index_field in self._created_indexes[table_name]):
                return
            
            index_name = f"idx_{table_name}_{index_field}"
            
            # 检查索引是否已存在
            cursor = await db.execute('''
                SELECT name FROM sqlite_master 
                WHERE type='index' AND name=?
            ''', (index_name,))
            index_exists = await cursor.fetchone()
            
            if not index_exists:
                await self._create_index(db, table_name, index_field)
            else:
                self.logger.info(f"索引 {index_name} 已存在，跳过创建")
                # 记录到缓存中
                if table_name not in self._created_indexes:
                    self._created_indexes[table_name] = {}
                self._created_indexes[table_name][index_field] = True
                
        except Exception as e:
            self.logger.error(f"检查/创建索引时发生错误: {e}")
            raise
    
    async def save(self, data: list, source_url: str, table_name: str, index_field: str = None):
        """
        保存数据到SQLite数据库
        
        Args:
            data: 解析后的数据列表
            source_url: 数据源URL
            table_name: 表名
            index_field: 需要创建索引的字段名（可选）
        """
        if not self._initialized:
            raise RuntimeError("SqliteSaver未正确初始化")
        
        if not self._sqlite_enabled:
            self.logger.warning("SQLite未启用，跳过保存")
            return
        
        if not data:
            self.logger.warning("没有数据需要保存")
            return
        if not source_url:
            self.logger.warning("没有数据来源URL")
            return
        if not table_name:
            self.logger.error("没有表名")
            return
        
        try:
            await self._save_to_sqlite(data, source_url, table_name, index_field)
        except Exception as e:
            self.logger.error(f"保存数据时发生错误: {e}")
            raise
    
    async def _save_to_sqlite(self, data: list, source_url: str, table_name: str, index_field: str = None):
        """保存数据到SQLite数据库"""
        try:
            # 确保数据库目录存在
            os.makedirs(os.path.dirname(self._db_path), exist_ok=True)
            
            self.logger.info(f"正在保存数据到SQLite数据库: {self._db_path} (数据库: {self._database_name})")
            
            async with aiosqlite.connect(self._db_path) as db:
                # 检查表是否存在
                cursor = await db.execute('''
                    SELECT name FROM sqlite_master WHERE type='table' AND name=?
                ''', (table_name,))
                table_exists = await cursor.fetchone()
                
                # 如果表已存在且指定了索引字段，检查是否需要创建索引
                if table_exists and index_field and self._validate_index_field(index_field, data):
                    await self._create_index_if_not_exists(db, table_name, index_field)
                
                if not table_exists and data:
                    # 根据data动态创建表结构
                    # 获取第一个数据项的所有键作为字段名
                    sample_item = data[0]
                    fields = list(sample_item.keys())
                    
                    # 构建CREATE TABLE语句
                    field_definitions = ['id INTEGER PRIMARY KEY AUTOINCREMENT']
                    field_definitions.extend([f'{field} TEXT' for field in fields])
                    field_definitions.append('source_url TEXT')
                    field_definitions.append('created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP')
                    
                    create_table_sql = f'''
                        CREATE TABLE {table_name} (
                            {', '.join(field_definitions)}
                        )
                    '''
                    
                    await db.execute(create_table_sql)
                    self.logger.info(f"动态创建数据表: {table_name}")
                    
                    # 创建索引（如果指定了索引字段且字段有效）
                    if self._validate_index_field(index_field, data):
                        await self._create_index(db, table_name, index_field)
                    elif index_field:
                        self.logger.warning(f"指定的索引字段 '{index_field}' 在数据中不存在，跳过索引创建")
                elif not data:
                    # 如果没有数据，报错
                    self.logger.error(f"没有数据需要保存")
                    return
                
                # 插入数据
                if data:
                    # 获取字段名
                    sample_item = data[0]
                    fields = list(sample_item.keys())
                    
                    # 构建INSERT语句
                    placeholders = ', '.join(['?' for _ in fields] + ['?'])  # 为source_url添加额外的占位符
                    field_names = ', '.join(fields + ['source_url'])
                    insert_sql = f'INSERT INTO {table_name} ({field_names}) VALUES ({placeholders})'
                    
                    # 准备批量插入的数据
                    batch_data = []
                    for item in data:
                        values = [item.get(field, '') for field in fields] + [source_url]
                        batch_data.append(values)
                    
                    # 使用executemany进行批量插入
                    await db.executemany(insert_sql, batch_data)
                
                await db.commit()
                self.logger.info(f"成功保存 {len(data)} 条数据到SQLite数据库: {self._db_path}")
                
        except sqlite3.OperationalError as e:
            error_msg = str(e).lower()
            if 'no such table' in error_msg:
                self.logger.error(f"数据表 {table_name} 不存在，且无法创建: {e}")
            elif 'database is locked' in error_msg:
                self.logger.error(f"数据库被锁定，请检查是否有其他进程正在使用: {e}")
            elif 'disk i/o error' in error_msg:
                self.logger.error(f"磁盘I/O错误，可能是磁盘空间不足或权限问题: {e}")
            else:
                self.logger.error(f"数据库操作错误: {e}")
            raise
            
        except sqlite3.IntegrityError as e:
            self.logger.error(f"数据完整性错误，可能是主键冲突或约束违反: {e}")
            raise
            
        except sqlite3.ProgrammingError as e:
            self.logger.error(f"SQL语法错误或编程错误: {e}")
            raise
            
        except sqlite3.DatabaseError as e:
            self.logger.error(f"数据库错误: {e}")
            raise
            
        except FileNotFoundError as e:
            self.logger.error(f"数据库文件路径错误: {e}")
            raise
            
        except PermissionError as e:
            self.logger.error(f"数据库文件权限错误: {e}")
            raise
            
        except OSError as e:
            if e.errno == 28:  # No space left on device
                self.logger.error(f"磁盘空间不足，无法保存数据: {e}")
            else:
                self.logger.error(f"系统错误: {e}")
            raise
            
        except Exception as e:
            self.logger.error(f"保存数据到SQLite时发生未知错误: {e}")
            raise
    
    def is_enabled(self) -> bool:
        """检查SQLite是否启用"""
        return self._sqlite_enabled
    
    def get_db_path(self) -> str:
        """获取数据库路径"""
        return self._db_path
    
    def get_database_name(self) -> str:
        """获取数据库名称"""
        return self._database_name
    
    def get_created_indexes(self) -> dict:
        """获取已创建索引的缓存信息"""
        return self._created_indexes.copy()
    
    def clear_index_cache(self):
        """清空索引缓存（用于重新加载配置或重置状态）"""
        self._created_indexes.clear()
        self.logger.info("索引缓存已清空")
    
    def is_index_cached(self, table_name: str, index_field: str) -> bool:
        """检查指定表的字段索引是否已缓存"""
        return (table_name in self._created_indexes and 
                index_field in self._created_indexes[table_name])